import numpy as np
import torch
from torch.autograd import Variable
from torchvision import datasets
from torchvision.transforms import transforms
from torchviz import make_dot


# 展示图片
def to_img(x):
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)
    out = out.view(-1, 1, 28, 28)
    return out


# 模型可视化，生成gv文件
def draw_model(model, input_size):
    input = torch.rand(input_size)
    output = model(input)
    g = make_dot(output)
    g.view()


# 加载数据集
def loadData(batch_size):
    # 图像预处理
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5], std=[0.5])
    ])
    # 获取MNIST数据集
    mnist = datasets.MNIST(
        root='./data/', train=True, transform=transform, download=False)
    # 加载MNIST数据集
    dataloader = torch.utils.data.DataLoader(
        dataset=mnist, batch_size=batch_size, shuffle=True)
    return dataloader


# 生成csv文件
def make_train_set(model, batch_size, noise_dim, dataloader, normal):
    img_list = []
    label_list = []
    Tensor = torch.FloatTensor
    for i in range(200):
        if normal == 0:
            noise = Variable(torch.randn(batch_size, noise_dim))    # 随机分布
        else:
            noise = Variable(Tensor(np.random.normal(0, 1, (batch_size, noise_dim))))  # 正态分布
        fake_img = model(noise)
        img_list.append(fake_img.view(-1).data.numpy())
        label_list.append(0)

    for i, (img, _) in enumerate(dataloader):
        if i == 200:
            break
        real_img = Variable(img)
        img_list.append(real_img.view(-1).data.numpy())
        label_list.append(1)
    return img_list, label_list
